パブリックプレビュー版のSnowpark ML Model Registryで、SnowflakeでのMLOpsのポイントを確認してみた

データアナリティクス事業本部 機械学習チームの鈴木です。

先日、Snowpark MLのModel Registryのパブリックプレビューが開始になりました。

この公開までの間、Snowpark MLでもいくつかの重要なアップデートがあり、それらが合わさってSnowflakeでのモデルのデプロイや管理がかなり使いやすくなったように思ったので、改めて触ってみました。

個人的には、

• 前処理も含めたパイプラインを管理することで、Snowflakeの特徴量用のマートテーブルを機械学習モデルで直接利用できるようになった
• モデルのバージョン管理が容易になった

という２点が大きなポイントだと感じています。こちらは記事の後半で紹介したいポイントとして記載します。

最初に

Model RegistryはSnowpark MLにおけるMLOps向けの機能の一つで、Snowflakeで機械学習モデルとそのメタデータを安全に管理することができます。

Model Registryがなかったときは、モデルをto_sklearnメソッドでscikit-learnのインスタンスに変換するなどしてUDFで実行できるようにし、UDFとしてデプロイするのがオーソドックスな戦略だったと思います。この場合、モデルのバージョン管理のためUDFの再作成時に数バージョン前のものは残しておくなど作り込みが必要でしたし、実装もそこそこ複雑でした。

また、Snowpark MLの古いバージョンだと前処理用の仕組みである各種変換やPipelineはscikit-learnのインスタンスに変換できなかったので、前処理部分は別途ストアドプロシージャとタスクで定期実行する必要があったかと思います。

この数ヶ月のアップデートでPipelineにもto_sklearnメソッドなどが追加されたほか、パブリックプレビュー版のModel Registryではパイプラインごとモデルをデプロイできるため、Snowpark MLで作成した機械学習モデルを運用していく上での便利さが段違いにパワーアップした認識です。

イメージとしては以下のブログが大変分かりやすいです。

Model Registryについては以下の公式ガイドとAPIリファレンスを確認して頂くと良いと思います。

本稿では、上記ブログ・ガイドを参考にしつつ、実際に触ってみての感想や、もう少し細かいところについて記載して、よりModel Registry（パブリックプレビュー版）の利用イメージが膨らめばと思っています。

前提と前準備

データ・コード・環境について

このブログでは、以下の『Intro to Machine Learning with Snowpark ML』の内容を題材に、Model Registryの便利さを確認します。

※ クイックスタートの更新によりリンクが切れている場合は、GETTING STARTED WITH SNOWFLAKEから検索してください。

データセットは、クイックスタートの内容にしたがってダイアモンドデータセットをDIAMONDSテーブルに用意しました。

Snowpark MLは1.2.3を使用しました。ローカル環境に構築したJupyter Notebook環境よりSnowparkでSnowflakeに接続しました。

なお、以降で紹介するコードはクイックスタートおよびガイドの内容を引用・参考にしています。

パイプラインの定義とトレーニング

登録済みのDIAMONDSテーブルから以下のようにパイプラインをトレーニングしました。

session作成などは少し長くなるため省略しています。詳しくは『クイックスタートでSnowpark ML Modelingを学んだのでポイントをご紹介』をご確認ください。

import snowflake.ml.modeling.preprocessing as snowml
from snowflake.ml.modeling.pipeline import Pipeline
from snowflake.ml.modeling.model_selection import GridSearchCV
from snowflake.ml.modeling.xgboost import XGBRegressor

# Snowflakeのテーブルからのデータの取得
# sessionはあらかじめ作成しておく
DEMO_TABLE = 'DIAMONDS'
input_tbl = f"{session.get_current_database()}.{session.get_current_schema()}.{DEMO_TABLE}"
diamonds_df = session.table(input_tbl)

# データセットの分割
diamonds_train_df, diamonds_test_df = diamonds_df.random_split(weights=[0.9, 0.1], seed=0)

# 前処理・学習・推論のためのカラム定義
CATEGORICAL_COLUMNS = ["CUT", "COLOR", "CLARITY"]
CATEGORICAL_COLUMNS_OE = ["CUT_OE", "COLOR_OE", "CLARITY_OE"]
NUMERICAL_COLUMNS = ["CARAT", "DEPTH", "TABLE_PCT", "X", "Y", "Z"]
categories = {
    "CUT": np.array(["IDEAL", "PREMIUM", "VERY_GOOD", "GOOD", "FAIR"]),
    "CLARITY": np.array(["IF", "VVS1", "VVS2", "VS1", "VS2", "SI1", "SI2", "I1", "I2", "I3"]),
    "COLOR": np.array(['D', 'E', 'F', 'G', 'H', 'I', 'J']),
}
LABEL_COLUMNS = ['PRICE']
OUTPUT_COLUMNS = ['PREDICTED_PRICE']

# 機械学習パイプラインの定義
pipeline = Pipeline(
    steps=[
            (
                "OE",
                snowml.OrdinalEncoder(
                    input_cols=CATEGORICAL_COLUMNS,
                    output_cols=CATEGORICAL_COLUMNS_OE,
                    categories=categories,
                )
            ),
            (
                "MMS",
                snowml.MinMaxScaler(
                    clip=True,
                    input_cols=NUMERICAL_COLUMNS,
                    output_cols=NUMERICAL_COLUMNS,
                )
            ),
            (
                "GridSearchCV",
                GridSearchCV(
                    estimator=XGBRegressor(random_state=42),
                        param_grid={
                                        "n_estimators":[300, 400],
                                        "learning_rate":[0.1, 0.2],
                                        },
                        n_jobs = -1,
                        scoring="neg_mean_squared_error",
                        input_cols=CATEGORICAL_COLUMNS_OE+NUMERICAL_COLUMNS,
                        label_cols=LABEL_COLUMNS,
                        output_cols=OUTPUT_COLUMNS
                )
            )
    ]
)

# モデルのトレーニング
pipeline.fit(diamonds_train_df)

紹介したいポイント

Model Registryや直近のアップデートにより、Snowpark MLがより使いやすくなったと思うポイントについて何点かピックアップしてご紹介します。

パイプラインのscikit-learnインスタンスへの変換

まず、Snowpark MLのv1.1.1以降のAPIリファレンスではPipelineおよび各種変換にto_sklearn()が実装されたことが明記されました。

以下のように、Scikit-learnのパイプラインに変換することができます。

pipeline.to_sklearn()

Snowpark MLのパイプラインの変換はscikit-learnと若干異なるのですが、to_sklearn()により適宜ColumnTransformerなどに置き換えられていることが分かります。

Model Registryへの登録

snowflake.ml.registry.Registryのlog_modelメソッドによりモデルやパイプラインをModel Registryに登録することができます。

from snowflake.ml.registry import Registry

reg = Registry(session=session)

reg.log_model(
    model_name="diamond_xgboost_regressor",
    version_name="V_1",
    model=pipeline,
    metrics=metrics
)

特にメトリクスは以下のように計算したものを登録しておくことができます。

# テストデータでの推論
result_df = pipeline.predict(diamonds_test_df)

# MSEの計算
mse = mean_squared_error(df=result_df, 
                         y_true_col_names="PRICE", 
                         y_pred_col_names="PREDICTED_PRICE")

# log_modelに渡すmetrics辞書の作成
metrics = {
    "MSE": mse
}

log_modelメソッドにはほかにもいくつかのメタデータを登録しておくことができます。

登録したモデルの管理

snowflake.ml.registry.Registryのshow_modelsメソッドで登録したモデルを確認することができました。

# 登録してあるモデルの取得
model_df = reg.show_models()

特に、新しいバージョンのモデルを登録した場合に、versionsが増えていくことが分かります。UDFなどと異なり、単に新しいものをデプロイしても以前のものが削除しされず、バージョンで指定できるのは運用が非常に簡単になるのでとても強力です。

SQLからも確認可能です。こちらはSHOW MODELSで表示できます。

モデルバージョンの利用

モデルのうちバージョンを指定し、推論やメタデータの確認を行うことができます。

ここでは推論の例を確認します。V_1のバージョンを使い、runでpredictメソッドを指定して推論を実行しました。

mv = m.version("V_1")
remote_prediction = mv.run(diamonds_test_df, function_name="predict")

Calling Model Methods in SQLに記載のように、SQLからも推論ができました。

-- 『Getting Started with Snowpark Model Registry』ブログを参考
SELECT *, 
       diamond_xgboost_regressor!PREDICT(*):PREDICTED_PRICE::FLOAT AS PREDICTED_PRICE 
FROM 
(
    SELECT * EXCLUDE PRICE 
    FROM DIAMONDS
    LIMIT 10
);

モデルのバージョンごとのメタデータの確認はPythonからも可能ですが、閲覧用途だとSQLが多いのかなと想像しています。こちらはSHOW VERSIONS IN MODELで表示できます。

SHOW VERSIONS IN MODEL DIAMOND_XGBOOST_REGRESSOR;

終わりに

今回は現在パブリックプレビュー中のSnowpark ML Model Registryについてご紹介しました。

ポイントとしてみた通り、パイプラインを作ることでテーブルのデータそのままをSnowflakeのリソースを使い高速に前処理を行うことができ、モデルの学習・推論を簡単なAPIで実行することができました。

これにより、Snowflake上のデータを使った機械学習モデルの開発と運用が、格段に導入・実践しやすくなったと思います。

一般提供開始を楽しみにしつつ、気になる方はぜひ触ってみて頂ければと思います。